# -*- coding:utf-8 -*-
# @Author: shenyuyu
# @Time: 2023/6/29 19:59
# @File: qu_3_自定义udf函数返回字典.py
import string

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StringType, IntegerType
from pyspark.sql import functions as F

if __name__ == '__main__':
    spark = SparkSession.builder.appName("a").master("local[*]").getOrCreate()
    sc = spark.sparkContext
    rdd1 = sc.parallelize([[1], [2], [3]])
    df = rdd1.toDF(["num"])
    df.createTempView("nums")


    def return_dict(data):
        return {"num": data, "letter": string.ascii_letters[data]}

    return_dict2 = spark.udf.register("return_dict1", return_dict, StructType().
                                      add("num", IntegerType()).add("letter", StringType()))

    spark.sql("select return_dict1(num) from nums").show()
    print(spark.sql("select return_dict1(num) from nums").first())
    print(type(spark.sql("select return_dict1(num) from nums").first()))
    print(type(spark.sql("select return_dict1(num) from nums").first()["return_dict1(num)"]["num"]))
